{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "635d8ebb",
   "metadata": {},
   "source": [
    "# Memory Using SQLite\n",
    "\n",
    "- Author: [Heesun Moon](https://github.com/MoonHeesun)\n",
    "- Peer Review: [harheem](https://github.com/harheem), [gyjong](https://github.com/gyjong)\n",
    "- Proofread : [Juni Lee](https://www.linkedin.com/in/ee-juni)\n",
    "- This is a part of [LangChain Open Tutorial](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial)\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/05-Memory/09-Memory-Using-SQLite.ipynb)[![Open in GitHub](https://img.shields.io/badge/Open%20in%20GitHub-181717?style=flat-square&logo=github&logoColor=white)](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/05-Memory/09-Memory-Using-SQLite.ipynb)\n",
    "## Overview\n",
    "\n",
    "This tutorial explains the ```SQLChatMessageHistory``` class, which allows storing chat history in any database supported by ```SQLAlchemy```.\n",
    "\n",
    "```Structured Query Language (SQL)``` is a domain-specific language used in programming and designed for managing data held in a Relational Database Management System (RDBMS), or for stream processing in a Relational Data Stream Management System (RDSMS). It is particularly useful for handling structured data, including relationships between entities and variables.\n",
    "\n",
    "```SQLAlchemy``` is an open-source **SQL** toolkit and Object-Relational Mapper (ORM) for the Python programming language, released under the MIT License.\n",
    "\n",
    "To use a database other than ```SQLite```, please make sure to install the appropriate database driver first.\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "- [Overview](#overview)\n",
    "- [Environment Setup](#environment-setup)\n",
    "- [Usage](#Usage)\n",
    "- [Chaining](#Chaining)\n",
    "\n",
    "### References\n",
    "\n",
    "- [Wikipedia: SQL](https://en.wikipedia.org/wiki/SQL)\n",
    "- [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy)\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6c7aba4",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Set up the environment. You may refer to [Environment Setup](https://wikidocs.net/257836) for more details.\n",
    "\n",
    "**[Note]**\n",
    "- ```langchain-opentutorial``` is a package that provides a set of easy-to-use environment setup, useful functions and utilities for tutorials. \n",
    "- You can checkout the [```langchain-opentutorial```](https://github.com/LangChain-OpenTutorial/langchain-opentutorial-pypi) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21943adb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install langchain_opentutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f25ec196",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages\n",
    "from langchain_opentutorial import package\n",
    "\n",
    "package.install(\n",
    "    [\n",
    "        \"langchain_community\",\n",
    "        \"langchain_openai\",\n",
    "        \"langchain_core\",\n",
    "        \"SQLAlchemy\",\n",
    "    ],\n",
    "    verbose=False,\n",
    "    upgrade=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b012261d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set environment variables\n",
    "from langchain_opentutorial import set_env\n",
    "\n",
    "set_env(\n",
    "    {\n",
    "        \"OPENAI_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_TRACING_V2\": \"true\",\n",
    "        \"LANGCHAIN_ENDPOINT\": \"https://api.smith.langchain.com\",\n",
    "        \"LANGCHAIN_PROJECT\": \"MemoryUsingSQLite\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ec228d5",
   "metadata": {},
   "source": [
    "You can alternatively set ```OPENAI_API_KEY``` in ```.env``` file and load it.\n",
    "\n",
    "[Note] This is not necessary if you've already set ```OPENAI_API_KEY``` in previous steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "763d9c57",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa00c3f4",
   "metadata": {},
   "source": [
    "## Usage\n",
    "\n",
    "To use the storage, you need to provide only the following 2 things:\n",
    "\n",
    "1. ```session_id``` - A unique identifier for the session, such as a user name, email, chat ID, etc.\n",
    "\n",
    "2. ```connection``` - A string that specifies the database connection. This string will be passed to SQLAlchemy's ```create_engine``` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "a466cc09",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.chat_message_histories import SQLChatMessageHistory\n",
    "\n",
    "# Initialize chat history with session ID and database connection.\n",
    "chat_message_history = SQLChatMessageHistory(\n",
    "    session_id=\"sql_history\", connection=\"sqlite:///sqlite.db\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "777c8d5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a user message\n",
    "chat_message_history.add_user_message(\n",
    "    \"Hello, nice to meet you! My name is Heesun :) I'm a LangChain developer. I look forward to working with you!\"\n",
    ")\n",
    "# Add an AI message\n",
    "chat_message_history.add_ai_message(\n",
    "    \"Hi, Heesun! Nice to meet you. I look forward to working with you too!\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bca8e6c7",
   "metadata": {},
   "source": [
    "Now, let's check the stored conversation history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "06fc3fac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content=\"Hello, nice to meet you! My name is Heesun :) I'm a LangChain developer. I look forward to working with you!\", additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='Hi, Heesun! Nice to meet you. I look forward to working with you too!', additional_kwargs={}, response_metadata={})]"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_message_history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d97c86b",
   "metadata": {},
   "source": [
    "You can also clear the session memory from db:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "b07e7029",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Clear the session memory\n",
    "chat_message_history.clear()\n",
    "chat_message_history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bccfeeb",
   "metadata": {},
   "source": [
    "### Adding Metadata\n",
    "\n",
    "**Metadata** can be added by directly creating ```HumanMessage``` and ```AIMessage``` objects. This approach enables flexible data handling and logging.\n",
    "\n",
    "**Parameters**:\n",
    "- ```additional_kwargs``` - Stores custom tags or metadata, such as priority or task type.\n",
    "\n",
    "- ```response_metadata``` - Captures AI response details, including model, timestamp, and token count.\n",
    "\n",
    "These fields enhance debugging and task tracking through detailed data storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "6783af0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "# Add a user message with additional metadata.\n",
    "user_message = HumanMessage(\n",
    "    content=\"Can you help me summarize this text?\",\n",
    "    additional_kwargs={\"task\": \"summarization\"},\n",
    ")\n",
    "\n",
    "# Add the message to chat history.\n",
    "chat_message_history.add_message(user_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "168c7129",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='Can you help me summarize this text?', additional_kwargs={'task': 'summarization'}, response_metadata={})]"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_message_history.messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "4b8cdacf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import AIMessage\n",
    "\n",
    "# Add an AI message with response metadata.\n",
    "ai_message = AIMessage(\n",
    "    content=\"Sure! Here's the summary of the provided text.\",\n",
    "    response_metadata={\"model\": \"gpt-4\", \"token_count\": 30, \"response_time\": \"150ms\"},\n",
    ")\n",
    "\n",
    "# Add the message to chat history.\n",
    "chat_message_history.add_message(ai_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "4beb2d42",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='Can you help me summarize this text?', additional_kwargs={'task': 'summarization'}, response_metadata={}),\n",
       " AIMessage(content=\"Sure! Here's the summary of the provided text.\", additional_kwargs={}, response_metadata={'model': 'gpt-4', 'token_count': 30, 'response_time': '150ms'})]"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_message_history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b9d7af2",
   "metadata": {},
   "source": [
    "## Chaining\n",
    "\n",
    "You can easily integrate this chat history class with [LCEL Runnables](https://wikidocs.net/235884)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "d318bc5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import (\n",
    "    ChatPromptTemplate,\n",
    "    MessagesPlaceholder,\n",
    ")\n",
    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_openai import ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "8f9e7b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant.\"),\n",
    "        # Placeholder for chat history\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Chaining\n",
    "chain = prompt | ChatOpenAI(model_name=\"gpt-4o\") | StrOutputParser()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c59c73b",
   "metadata": {},
   "source": [
    "The following shows how to create a function that returns chat history from ```sqlite.db```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "eaf1d642",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_chat_history(user_id, conversation_id):\n",
    "    return SQLChatMessageHistory(\n",
    "        table_name=user_id,\n",
    "        session_id=conversation_id,\n",
    "        connection=\"sqlite:///sqlite.db\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "244eedab",
   "metadata": {},
   "source": [
    "Set ```config_fields``` to provide reference information when retrieving conversation details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "3c5d4580",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables.utils import ConfigurableFieldSpec\n",
    "\n",
    "config_fields = [\n",
    "    ConfigurableFieldSpec(\n",
    "        id=\"user_id\",\n",
    "        annotation=str,\n",
    "        name=\"User ID\",\n",
    "        description=\"Unique identifier for a user.\",\n",
    "        default=\"\",\n",
    "        is_shared=True,\n",
    "    ),\n",
    "    ConfigurableFieldSpec(\n",
    "        id=\"conversation_id\",\n",
    "        annotation=str,\n",
    "        name=\"Conversation ID\",\n",
    "        description=\"Unique identifier for a conversation.\",\n",
    "        default=\"\",\n",
    "        is_shared=True,\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "eb9f0967",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain_with_history = RunnableWithMessageHistory(\n",
    "    chain,\n",
    "    get_chat_history,\n",
    "    input_messages_key=\"question\",\n",
    "    history_messages_key=\"chat_history\",\n",
    "    # Set parameters for retrieving chat history\n",
    "    history_factory_config=config_fields,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "953c4c5a",
   "metadata": {},
   "source": [
    "Set the ```\"user_id\"``` and ```\"conversation_id\"``` key-value pairs under the ```\"configurable\"``` key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "2b7e4455",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Config settings\n",
    "config = {\"configurable\": {\"user_id\": \"user1\", \"conversation_id\": \"conversation1\"}}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d76e860",
   "metadata": {},
   "source": [
    "Let's ask a question about the name. If there is any previously saved conversation history, it will provide the correct response.  \n",
    "\n",
    "- Use the ```invoke``` method of the ```chain_with_history``` object to generate an answer to the question.  \n",
    "- Pass a question dictionary and ```config``` settings to the ```invoke``` method as inputs.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "15ac9974",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Hi Heesun! Nice to meet you again. How can I help you today?'"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Execute by passing the question and config\n",
    "chain_with_history.invoke(\n",
    "    {\"question\": \"Hi, nice to meet you. My name is Heesun.\"}, config\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "ea24385c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Your name is Heesun.'"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Execute a follow-up question\n",
    "chain_with_history.invoke({\"question\": \"What is my name?\"}, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49e43ba3",
   "metadata": {},
   "source": [
    "This time, set the same ```user_id``` but use a different value for ```conversation_id```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "da5a49c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"I'm sorry, but I don't have access to personal information, so I don't know your name. If you'd like, you can tell me your name, and I can address you by it.\""
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Config settings\n",
    "config = {\"configurable\": {\"user_id\": \"user1\", \"conversation_id\": \"conversation2\"}}\n",
    "\n",
    "# Execute by passing the question and config\n",
    "chain_with_history.invoke({\"question\": \"What is my name?\"}, config)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
